# 实现Test1
from dolfinx import mesh, plot, fem, default_real_type
from dolfinx.fem import functionspace
from mpi4py import MPI
import pyvista as pv
import numpy as np
import basix
from basix.ufl import element, mixed_element
import ufl

# 算例

# 固定压力点
def marker(x): return np.isclose(x[0], 1) & np.isclose(x[1], 1)

def true_velocity_function(x):
    """
    速度真解 函数
    """
    u1 = -4 * x[1] * (4 - x[0]**2)**2 * (4 - x[1]**2)
    u2 = 4 * x[0] * (4 - x[0]**2) * (4 - x[1]**2)**2
    return [u1, u2]


def true_pressure_function(x):
    """
    压强真解 函数
    """
    return 150 * np.sin(x[0])


def lambda_function(x):
    """拉格朗日乘子真解

    Args:
        x (_type_): msh
    """
    return [np.exp(x[0]), np.exp(x[1])]


def displacement_function(x):
    """位移真解
    """
    return true_velocity_function(x)  # 位移真解和速度真解相同


def c(mu, Y):
    """双线性形式 c(μ, Y)
    """
    return (ufl.inner(ufl.sym(ufl.grad(mu)), ufl.sym(ufl.grad(Y))) + ufl.inner(mu, Y)) * dx_solid


def identity_map(s):
    return s  # s 为固体参考域坐标 (x, y)


def fluid_boundary(x):
    """流体边界"""
    return np.isclose(x[0], 2.0) | np.isclose(
        x[0], -2.0) | np.isclose(x[1], -2.0) | np.isclose(x[1], 2.0)  

def solid_boundary(x):
    """固体边界"""
    return np.isclose(x[0], 1.0) | np.isclose(
        x[0], -1.0) | np.isclose(x[1], -1.0) | np.isclose(x[1], 1.0)  


h_solid = 0.125  # 固体网格大小
h_fluid = 0.125
Lx_solid, Ly_solid = 2, 2  # 固体网格大小
Lx_fluid, Ly_fluid = 4, 4  # 流体网格大小
n_solid = [int(Lx_solid / h_solid), int(Ly_solid / h_solid)]
n_fluid = [int(Lx_fluid / h_fluid), int(Ly_fluid / h_fluid)]

solid_msh = mesh.create_rectangle(
    MPI.COMM_WORLD, [[-1, -1], [1, 1]], n=n_solid, cell_type=mesh.CellType.triangle)
fluid_msh = mesh.create_rectangle(
    MPI.COMM_WORLD, [[-2, -2], [2, 2]], n=n_fluid, cell_type=mesh.CellType.triangle)

# 定义有限元空间 https://defelement.org/elements/p1-iso-p2.html
P2_iso_ = basix.create_element(
    basix.ElementFamily.iso, basix.CellType.triangle, 1)  # type=FiniteElement
P2_iso_u = element(P2_iso_.family, fluid_msh.basix_cell(), degree=1,
                   shape=(fluid_msh.geometry.dim,), dtype=default_real_type)
P1_p = element("Lagrange", fluid_msh.basix_cell(), degree=1,
               dtype=default_real_type)
mel_fluid = mixed_element([P2_iso_u, P1_p])  # 定义在\Omega上速度-压力混合元
ME_fluid = functionspace(fluid_msh, mel_fluid)  # 实例化混合元空间


# 拉格朗日乘子 and 位移 耦合
P1_lagrange = element(
    "Lagrange", solid_msh.basix_cell(), degree=1, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)  # 拉格朗日乘子
P1_S = element(
    "Lagrange", solid_msh.basix_cell(), degree=1, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)  # 位移
mel_solid = mixed_element([P1_lagrange, P1_S])  # 定义在参考域上 拉格朗日乘子-位移混合元
ME_solid = functionspace(solid_msh, mel_solid)  # 实例化混合元空间
dx_fluid = ufl.Measure("dx", domain=fluid_msh)
dx_solid = ufl.Measure("dx", domain=solid_msh)


VT, _ = ME_fluid.sub(0).collapse()
PT, _ = ME_fluid.sub(1).collapse()
u_bc = fem.Function(VT)  # 真实速度函数空间
p_bc = fem.Function(PT)  # 真实压力函数空间
u_bc.interpolate(true_velocity_function)
p_bc.interpolate(true_pressure_function)

X_bar = ufl.SpatialCoordinate(solid_msh)  # 定义空间坐标
# 流体域测试函数和试函数
(v, q) = ufl.TestFunctions(ME_fluid)
(u, p) = ufl.TrialFunctions(ME_fluid)

# 固体域测试函数和试函数
(mu, Y) = ufl.TestFunctions(ME_solid)
(lambda_, X) = ufl.TrialFunctions(ME_solid)
# 流体域体力项 f（包含 -Δu + ∇p + c(λ, v)）
# 这里需要实现第五章的算法
# f_fluid = -ufl.div(ufl.grad(u)) + ufl.grad(p) + c(lambda_, v)
 
# 固体域体力项 g（包含 -ΔX - λ + g_∂B）
n_solid = ufl.FacetNormal(solid_msh)  # 固体域 单位外法向量
g_boundary = ufl.dot(ufl.grad(X), n_solid)
g_solid = -ufl.div(ufl.grad(X)) - lambda_ + g_boundary  # g_boundary 需提前定义

# 界面位移约束 d（u(X̄) - X）
# d = u(X_bar) - X
d = ufl.as_vector([0, 0])
 

# 设置流体边界自由度
fdim = fluid_msh.topology.dim - 1
facets = mesh.locate_entities_boundary(fluid_msh, fdim, fluid_boundary)
v_dofs = fem.locate_dofs_topological((ME_fluid.sub(0), VT), fdim, facets)
p_dofs = fem.locate_dofs_geometrical((ME_fluid.sub(1), PT), marker)
bc_vel = fem.dirichletbc(u_bc, v_dofs, ME_fluid.sub(0))  # 速度边界条件
bc_pressure = fem.dirichletbc(p_bc, p_dofs, ME_fluid.sub(1))  # 压力边界条件
bcs = [bc_vel, bc_pressure]  # 边界条件列表

# 双线性形 这里ufl.sym()的使用存在困惑 是否有必要?
a_f = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx_fluid  # a_f(u, v)
a_s = ufl.inner(ufl.sym(ufl.grad(X)), ufl.sym(
    ufl.grad(Y))) * dx_solid  # a_s(X, Y)
# c(mu, Y)形式 使用函数c(mu, Y)来表示双线性形式
# c（mu, v)形式 Chapter5 给出计算方法
# c1 = c(lambda_, v)  # c(λ, v)
c2 = -c(lambda_, Y) # c(λ, Y)
c3 = c(mu, d) # c(μ, u - X)


# 流体域线性形式
L_f = ufl.inner(f_fluid, v) * dx_fluid

# 固体域线性形式
L_s = ufl.inner(g_solid, Y) * dx_solid + ufl.inner(d, mu) * dx_solid


# 全局双线性形式 
a = a_f + a_s + c1 + c2 + c3
a -= ufl.div(v) * p * dx_fluid  # 流体域的压力项 
a -= ufl.div(u) * q * dx_fluid  # 质量守恒
# 全局线性形式
L = L_f + L_s

# 组装矩阵和向量
A = fem.petsc.assemble_matrix(a, bcs=bcs)
A.assemble()
b = fem.petsc.assemble_vector(L)
fem.petsc.apply_lifting(b, [a], [bcs])
b.ghostUpdate(addv=fem.petsc.ScatterMode.ADD,
              mode=fem.petsc.ScatterMode.REVERSE)
fem.petsc.set_bc(b, bcs)

# 创建并配置Krylov子空间求解器
solver = fem.petsc.LinearProblem(
    a, L, bcs=bcs,
    petsc_options={"ksp_type": "gmres", "pc_type": "lu"}
)

# 求解问题
solution = solver.solve()
